import cv2
import numpy as np
from skimage.metrics import structural_similarity as ssim
import os
import logging
import base64
from utils import convert_numpy_types

logger = logging.getLogger("image_similarity")


def compute_structural_similarity(img1_path, img2_path, threshold=0.8, params=None):
    """
    计算两个图像之间的结构相似性

    参数:
        img1_path: 第一张图像的路径
        img2_path: 第二张图像的路径
        threshold: 结构相似性阈值，默认0.8
        params: 其他参数字典，可以包含:
            - win_size: 窗口大小 (默认7)
            - sigma: 高斯滤波器标准差 (默认1.5)
            - alpha: 亮度权重 (默认0.0)
            - beta: 对比度权重 (默认0.0)
            - gamma: 结构权重 (默认1.0)

    返回:
        包含相似度分数和差异图的字典
    """
    # 设置默认参数
    if params is None:
        params = {}

    win_size = params.get('win_size', 7)
    sigma = params.get('sigma', 1.5)
    alpha = params.get('alpha', 0.0)  # 亮度权重
    beta = params.get('beta', 0.0)    # 对比度权重
    gamma = params.get('gamma', 1.0)  # 结构权重

    try:
        # 读取图像
        img1 = cv2.imread(img1_path)
        img2 = cv2.imread(img2_path)

        if img1 is None or img2 is None:
            logger.error(f"无法读取图像文件: {img1_path} 或 {img2_path}")
            return {
                "error": "无法读取图像文件",
                "structure_score": 0.0,
                "match": False
            }, 400

        logger.info(f"图像1尺寸: {img1.shape}, 图像2尺寸: {img2.shape}")

        # 检查图像尺寸
        if img1.shape != img2.shape:
            logger.warning(f"图像尺寸不同: {img1.shape} vs {img2.shape}")
            # 可选: 调整尺寸使两图相同
            img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))

        # 转换为RGB和灰度
        img1_rgb = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
        img2_rgb = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
        gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
        gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

        # 计算SSIM
        structure_score, structure_diff = ssim(
            gray1, gray2,
            full=True,
            data_range=255,
            gaussian_weights=True,
            sigma=sigma,
            use_sample_covariance=False,
            win_size=win_size,
            K1=0.01, K2=0.03,
            alpha=alpha,
            beta=beta,
            gamma=gamma
        )

        # 处理差异图
        structure_diff = (structure_diff * 255).astype("uint8")
        _, thresh = cv2.threshold(
            structure_diff, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)

        # 标记差异
        contours, _ = cv2.findContours(
            thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        diff_highlighted = img1_rgb.copy()
        cv2.drawContours(diff_highlighted, contours, -1, (0, 0, 255), 2)

        # 转换差异图为Base64编码
        _, diff_buffer = cv2.imencode('.png', structure_diff)
        diff_base64 = base64.b64encode(diff_buffer).decode('utf-8')

        _, highlighted_buffer = cv2.imencode(
            '.png', cv2.cvtColor(diff_highlighted, cv2.COLOR_RGB2BGR))
        highlighted_base64 = base64.b64encode(
            highlighted_buffer).decode('utf-8')

        # 结果
        result = {
            'structure_score': float(structure_score),
            'match': structure_score >= threshold,
            'diff_image_base64': highlighted_base64,
            'diff_map_base64': diff_base64,
            'threshold': threshold
        }

        # 转换所有NumPy类型
        result = convert_numpy_types(result)

        logger.info(
            f"结构相似度分数: {result['structure_score']:.4f}, 阈值: {result['threshold']}")
        return result, 200

    except Exception as e:
        logger.exception(f"计算结构相似度时出错: {str(e)}")
        return {"error": f"处理过程中出错: {str(e)}"}, 500
